import io
from http.cookies import SimpleCookie
from typing import Annotated

import aiohttp
from cachetools import TTLCache
from fastapi import APIRouter, Header, Request, Response
from starlette.responses import StreamingResponse

from app_v2.bili_api.api import ApiSession
from app_v2.bili_api.common import headers_
from app_v2.utils.RedicManager import RedisPool

code_cache = TTLCache(maxsize=10, ttl=60)
proxy = APIRouter(tags=["proxy"])


@proxy.get("/qrcode")
async def qrcode():
    async with aiohttp.ClientSession() as session:
        async with session.get(
                "https://passport.bilibili.com/x/passport-login/web/qrcode/generate",
                headers=headers_) as resp:
            resp.raise_for_status()
            data = await resp.json()
        return data


# TODO: depends code_cache
@proxy.get("/qrcode/poll")
async def qrcode_poll(qrcode_key: str):
    async with aiohttp.ClientSession() as session:
        async with session.get(
                "https://passport.bilibili.com/x/passport-login/web/qrcode/poll",
                params={"qrcode_key": qrcode_key},
                headers=headers_
        ) as resp:
            resp.raise_for_status()
            data = await resp.json()
        if data.get("code") != 0:
            return {
                "code": 503,
                "message": "bilibili routes error"
            }
        if data["data"]["code"] == 0:
            simple = SimpleCookie()
            cookies = resp.headers.getall("Set-Cookie")
            for cookie in cookies:
                simple.load(cookie)
            cookies = "; ".join(f"{k}={v.value}" for k, v in simple.items())
            print(cookies)
            code_cache[qrcode_key] = cookies
        return data


@proxy.get("/search")
async def search(keyword: str, token: Annotated[str | None, Header()]):
    async with ApiSession(token) as session:
        return await session.search(keyword)


# faster but smaller
img_cache = TTLCache(maxsize=40, ttl=180)


@proxy.get("/proxy/cache/img")
async def proxy_img(url: str):
    print(f"PROXY {url}")
    if data := img_cache.get(url):
        print("CACHED HIT")
        img = io.BytesIO(data)
        return StreamingResponse(img, media_type="image/jpeg")

    async with aiohttp.ClientSession() as session:
        async with session.get(url, headers=headers_) as resp:
            resp.raise_for_status()
            img_cache[url] = await resp.content.read()
        img = io.BytesIO(img_cache[url])
        return StreamingResponse(img, media_type="image/jpeg")


@proxy.get("/proxy/img")
async def proxy_img(url: str, request: Request, response: Response):
    response.headers["Cache-Control"] = "public, max-age=3600, immutable"

    redis: RedisPool = request.app.state.redis_pool
    cache_key = redis.generate_key("img_cache", url)
    image_data = await redis.get_client("img_cache").get(cache_key)
    if image_data is not None:
        print("CACHED HIT")
        img = io.BytesIO(image_data)
        return StreamingResponse(img, media_type="image/jpeg")
    async with aiohttp.ClientSession() as session:
        async with session.get(url, headers=headers_) as resp:
            resp.raise_for_status()
            image_data = await resp.read()
    await redis.get_client("img_cache").setex(
        cache_key,
        3600 * 24,
        image_data
    )
    img = io.BytesIO(image_data)
    return StreamingResponse(img, media_type="image/jpeg")
